from pdf2image import convert_from_path
from PIL import Image
from pathlib import Path
import urllib.request
import json
import base64
import os
import shutil

__ocr_url = "http://lingkun.open.woa.com/youtu/ocrapi/advertise_ocr"


def pdf_to_png(pdf_path, image_path):
    images = convert_from_path(pdf_path, dpi=300)
    for i, image in enumerate(images):
        image.save(image_path + f"/{i}.png", 'PNG')


def split_image(image_dir, split_image_dir):
    path = Path(image_dir)
    for file_path in path.rglob('*.png'):
        # 打开图片
        image = Image.open(str(file_path))

        # 获取图片的宽度和高度
        width, height = image.size

        # 计算裁剪点（中间位置）
        left_width = width // 2
        num = int(file_path.name.split(".")[0])
        if (num % 2):
            left_width += 20
        else:
            left_width -= 15

        # 裁剪左边部分
        left_image = image.crop((0, 0, left_width, height))

        # 裁剪右边部分
        right_image = image.crop((left_width, 0, width, height))

        # 保存裁剪后的图片到指定目录
        left_image_path = f"{split_image_dir}/{file_path.name.split('.')[0]}-1.{file_path.name.split('.')[-1]}"
        right_image_path = f"{split_image_dir}/{file_path.name.split('.')[0]}-2.{file_path.name.split('.')[-1]}"
        left_image.save(left_image_path)
        right_image.save(right_image_path)


def get_ocr_text(image_dir, split_text_dir):
    path = Path(image_dir)
    for file_path in path.rglob('*.png'):
        with open(str(file_path), mode="rb") as fs:
            file_content = fs.read()
        body = {
            "session_id": "123",
            "app_id": "123",
            "image": base64.b64encode(file_content).decode("utf-8")
        }
        data = json.dumps(body).encode("utf-8")
        req = urllib.request.Request(url=__ocr_url, data=data, method="POST")
        response = urllib.request.urlopen(req)
        resp_json = response.read().decode("utf-8")
        resp = json.loads(resp_json)
        if resp["errorcode"] != 0:
            break
        text_line = []
        items = resp["items"]
        for item in items:
            text_line.append(item["itemstring"])

        if len(text_line) > 0 and ("新传说" in text_line[0] or "大千世界" in text_line[0]):
            text_line.pop(0)
        if len(text_line) > 0 and "10月上半月刊" in text_line[-1]:
            text_line.pop(len(text_line) - 1)
        if len(text_line) > 0 and text_line[-1].isdigit():
            text_line.pop(len(text_line) - 1)
        if len(text_line) > 0 and "题图、插图" in text_line[-1]:
            text_line.pop(len(text_line) - 1)
        if len(text_line) > 0 and "发稿编辑" in text_line[-1]:
            text_line.pop(len(text_line) - 1)


        text = "".join(text_line)

        with open(f"{split_text_dir}/{file_path.name.split('.')[0]}.txt", mode="w") as fs:
            fs.write(text)


def init_dir(dirs):
    for dir in dirs:
        if os.path.exists(dir):
            shutil.rmtree(dir)
        os.makedirs(dir, exist_ok=True)


def merge_file(input_files, output_file):
    with open(output_file, 'w', encoding='utf-8') as out_f:
        for input_file in input_files:
            with open(input_file, mode='r') as input_f:
                out_f.write(input_f.read())


if __name__ == "__main__":
    root_path = "/Users/chenbingkang/Documents/数据/故事数据/故事会-2018-19/"
    pdf_path = root_path + "故事会-2018-19.pdf"
    merge_list = [9, 13, 18, 22, 26, 29, 32, 36, 39, 43, 48, 51, 53, 56, 61, 63, 67, 82, 83, 87, 89]

    image_dir = root_path + 'image'
    split_image_dir = root_path + "split_image"
    split_text_dir = root_path + "split_text"
    text_dir = root_path + "text"

    init_dir([image_dir, split_image_dir, split_text_dir, text_dir])
    pdf_to_png(pdf_path, image_dir)
    split_image(image_dir, split_image_dir)
    get_ocr_text(split_image_dir, split_text_dir)

    for i in range(len(merge_list) - 1):
        input_files = []
        for j in range(merge_list[i], merge_list[i + 1]):
            input_files.append(split_text_dir + f"/{j}-1.txt")
            input_files.append(split_text_dir + f"/{j}-2.txt")
        merge_file(input_files, text_dir + f"/{merge_list[i]}.txt")
